from scipy.interpolate import RBFInterpolator
from scipy.stats.qmc import Halton
from scipy.spatial import cKDTree
from scipy.interpolate import griddata
import numpy as np
import open3d as o3d
import math

def thick_filter(raw, r, iter):    
    polar = raw
    for i in range(iter):        
        tree = cKDTree(polar[:, 0:2])
        idx = tree.query_ball_point(polar[:, 0:2], r=r)
        polar[:, 2] = np.array([polar[idx[j], 2].sum() / len(idx[j]) for j in range(len(idx))])

    return polar[:, 2]

class PolarCalibrator:
    def __init__(self) -> None:
        self.__kdtree = None
        self.__calibration = None
        # self.__points = None
        # self.__values = None

    def train(self, data, ground_truth):
        self.__calibration = np.zeros((data.shape[0], data.shape[1] + 1), dtype=np.float64)
        self.__calibration[:, 0:-1] = data
        self.__calibration[:, -1] = ground_truth / data[:, -2]
        self.__kdtree = cKDTree(self.__calibration[:, 0:-1])
        # self.__points = data
        # self.__values = ground_truth / data[:, -2]

    def calibrate(self, data):
        d, i = self.__kdtree.query(data, k=8)
        e = (self.__calibration[i, -1] * d).sum(1) / d.sum(1)
        return data[:, -2] * e
        # return griddata(self.__points, self.__values, data, fill_value=0) * data[:, -2]

def cut(m):
    m = m.copy()
    m = m[:, [1, 2, 3, 5]]
    m[:, 3] = np.rad2deg(np.arccos(m[:, 3]))
    if np.min(m[:, 3]) < 0:
        print("")
    return m


def normalization(m):
    m = m.copy()
    b = [360, 180, 10000, 90]
    for i in range(4):
        m[:, i] = m[:, i] / b[i]
    return m


def unormalization(m):
    m = m.copy()
    b = [360, 180, 10000, 90]
    for i in range(4):
        m[:, i] = m[:, i] * b[i]
    return m

def calibrate(y1, y2):
    # k = 10
    # iter = 1
    # y0 = y2.copy()
    # y1[:, 2] = np.sqrt(np.square(thick_filter(transformSPH2XYZ(y1[:, 0:3]), k, iter)).sum(1))    
    # y2[:, 2] = np.sqrt(np.square(thick_filter(transformSPH2XYZ(y2[:, 0:3]), k, iter)).sum(1))
    y2[:, 3] = thick_filter(y2[:, 1:4], 0.001, 1)
    x1 = handle_x(y1)
    x2 = handle_x(y2)
    x1 = cut(x1)
    x2 = cut(x2)
    y1 = cut(y1)
    y2 = cut(y2) 
   
    get_pcd(y2)
    x1 = normalization(x1)
    y1 = normalization(y1)
    x2 = normalization(x2)
    y2 = normalization(y2)

    calib = PolarCalibrator()
    calib.train(x2[:10:, :], y2[:10:, 2])
    x3 = x2.copy()
    x3[:, 2] = calib.calibrate(x3)
    x2 = unormalization(x2)
    x3 = unormalization(x3)
    y2 = unormalization(y2)
    get_pcd(x3, y2)


def rbf(y1, y2):
    x1 = handle_x(y1)
    x2 = handle_x(y2)

    x1 = cut(x1)
    x2 = cut(x2)
    y1 = cut(y1)
    y2 = cut(y2)
    get_pcd(x2, y2)
    x1 = normalization(x1)
    y1 = normalization(y1)
    x2 = normalization(x2)
    y2 = normalization(y2)

    # x1 = unormalization(x1)
    # y1 = unormalization(y1)
    # x2 = unormalization(x2)
    # y2 = unormalization(y2)

    # get_pcd(x1, y1)
    # get_pcd(x2, y2)
    # 温度，水平转角，竖直转角，距离，反射率，cos

    # ‘linear’: -r
    # ‘thin_plate_spline’: r**2 * log(r)
    # ‘cubic’: r**3
    # ‘quintic’: -r**5
    # ‘multiquadric’: -sqrt(1 + r**2)
    # ‘inverse_multiquadric’: 1/sqrt(1 + r**2)
    # ‘inverse_quadratic’: 1/(1 + r**2)
    # ‘gaussian’: exp(-r**2)

    interpolation = RBFInterpolator(
        x2[:10:, :], y2[:10:, 2], neighbors=10, smoothing=0, kernel='linear')
    x3 = x2.copy()
    x3[:, 2] = interpolation(x3)
    x2 = unormalization(x2)
    x3 = unormalization(x3)
    get_pcd(x2, x3)


def handle_x(x1):
    y1 = x1.copy()
    acos = np.arccos(y1[:, 5])
    y1[:, 3] = y1[:, 3] * (1 + np.sin(acos * 10) / 10)
    return y1


def get_pcd2(m):
    m = transformSPH2XYZ(m)
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(m[:, [0, 1, 2]])
    pcd.paint_uniform_color([1, 0, 0])
    o3d.visualization.draw_geometries([pcd])


def get_pcd(m, n=None):
    pcds = []
    m = transformSPH2XYZ(m)    
    pcd1 = o3d.geometry.PointCloud()
    pcd1.points = o3d.utility.Vector3dVector(m[:, [0, 1, 2]])
    pcd1.paint_uniform_color([1, 0, 0])
    pcds.append(pcd1)
    # pcd1.colors = get_colors(m[:, 4])
    if n is not None:
        n = transformSPH2XYZ(n)
        pcd2 = o3d.geometry.PointCloud()
        pcd2.points = o3d.utility.Vector3dVector(n[:, [0, 1, 2]])
    # pcd2.colors = get_colors(m[:, 4])
        pcd2.paint_uniform_color([0, 0.9, 0])
        pcds.append(pcd2)
    
    o3d.visualization.draw_geometries(pcds)


def transformSPH2XYZ(_row):
    row = _row.copy()
    a = np.deg2rad(row[:, 0])  # 水平转角
    b = np.deg2rad(row[:, 1])  # 竖直转角
    r = row[:, 2]  # 距离
    row[:, 0] = r * np.sin(b) * np.cos(a)
    row[:, 1] = -r * np.sin(b) * np.sin(a)
    row[:, 2] = r * np.cos(b)
    return row

def get_colors(reflectivity_list):
    max_val = max(reflectivity_list)
    colors = []
    for row in reflectivity_list:
        c = row / max_val
        colors.append([c, c, c])
    return o3d.utility.Vector3dVector(colors)


if __name__ == '__main__':
    x1 = np.load("./x1.npy")
    y1 = np.load("./y1.npy")
    x2 = np.load("./x2.npy")
    # rbf(y1, x2)
    calibrate(y1, x2)
